import config
import paramiko
from scp import SCPClient
import os
import json
from openpyxl import Workbook, load_workbook
from config_reader import Config

# 读取测试配置信息
config_file = './config/test-config.yaml'
config = Config(config_file)
# 测试机ssh连接信息
ip = config.get('testclient', 'ip')
port = config.get('testclient', 'port')
username = config.get('testclient', 'username')
password = config.get('testclient', 'password')
# 实例GPU规格信息
gpu_type = config.get('gpuinfo', 'gpu_type') + "x" + config.get('gpuinfo', 'gpu_num')
# llmperf环境路径
env_activate_path = config.get('llmperfenv', 'env_activate_path')
llmperf_path = config.get('llmperfenv', 'llmperf_path')
# openai key和url
openai_api_key = config.get('openaienv', 'openai_api_key')
openai_api_base = config.get('openaienv', 'openai_api_base')


class SSHClient:
    def __init__(self, hostname, port, username, password):
        self.hostname = hostname
        self.port = port
        self.username = username
        self.password = password
        self.client = None

    def connect(self):
        self.client = paramiko.SSHClient()
        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.client.connect(self.hostname, self.port, self.username, self.password)
        self.shell = self.client.invoke_shell()
        print("SSH连接成功并打开了一个shell会话")

    def execute_command(self, command, timeout=60):
        if not self.shell:
            raise Exception("Shell会话未初始化")

        self.shell.send(command + '\n')
        output = ""
        while True:
            if self.shell.recv_ready():
                recv = self.shell.recv(1024).decode()
                output += recv
                if output.strip().endswith('$') or output.strip().endswith('#'):
                    break
        return output

    def is_directory(self, remote_path):
        sftp = self.client.open_sftp()
        try:
            return paramiko.SFTPAttributes.from_stat(sftp.stat(remote_path)).st_mode & 0o040000 == 0o040000
        except IOError:
            return False
        finally:
            sftp.close()

    def download_file(self, remote_path, local_path):
        with SCPClient(self.client.get_transport()) as scp:
            scp.get(remote_path, local_path)
            print(f"文件下载成功：{remote_path} 到 {local_path}")

    def download_directory(self, remote_path, local_path):
        with SCPClient(self.client.get_transport()) as scp:
            scp.get(remote_path, local_path, recursive=True)
            print(f"目录下载成功：{remote_path} 到 {local_path}")

    def close(self):
        if self.client:
            self.client.close()
            print("SSH连接关闭")


class AutomatedTest:
    def __init__(self, ssh_client, test_commands, remote_path, local_path):
        self.ssh_client = ssh_client
        self.test_commands = test_commands
        self.remote_path = remote_path
        self.local_path = local_path

    def run_test(self):
        try:
            # 连接SSH
            self.ssh_client.connect()

            # 执行测试命令
            for command in self.test_commands:
                print(f"执行命令： {command}")
                # output, error = self.ssh_client.execute_command(command, timeout=3600)
                output = self.ssh_client.execute_command(command, timeout=3600)
                print(f"命令输出: \n{output}")

            # 判断远程路径是文件还是目录并下载
            if self.ssh_client.is_directory(self.remote_path):
                print(f"下载目录: {self.remote_path} 到 {self.local_path}")
                self.ssh_client.download_directory(self.remote_path, self.local_path)
            else:
                print(f"下载文件: {self.remote_path} 到 {self.local_path}")
                self.ssh_client.download_file(self.remote_path, self.local_path)

        except Exception as e:
            print(f"发生错误： {e}")

        finally:
            self.ssh_client.close()


class LlmApiPerfTest():
    def __init__(self, *args, **kwargs):

        # 定义默认值
        defaults = {
            'model_name': 'baichuan2-13B-Chat',
            'mean_input_tokens': '550',
            'stddev_input_tokens': '150',
            'mean_output_tokens': '150',
            'stddev_output_tokens': '20',
            'max_num_completed_requests': '30',
            'timeout': '3600',
            'num_concurrent_requests': '1'
        }

        # 更新默认值
        defaults.update(kwargs)

        # 参数赋值给实例变量
        self.model_name = defaults['model_name']
        self.mean_input_tokens = defaults['mean_input_tokens']
        self.stddev_input_tokens = defaults['stddev_input_tokens']
        self.mean_output_tokens = defaults['mean_output_tokens']
        self.stddev_output_tokens = defaults['stddev_output_tokens']
        self.max_num_completed_requests = defaults['max_num_completed_requests']
        self.timeout = defaults['timeout']
        self.num_concurrent_requests = defaults['num_concurrent_requests']

        # 远程测试结果文件路径
        self.remote_path = "%s/result_outputs/%s_%s_%s_summary.json" % (
            llmperf_path, self.model_name, self.mean_input_tokens, self.mean_output_tokens)
        # 下载到本地的文件路径
        if not os.path.exists('./results'):
            os.makedirs('./results')
        self.local_path = "./results/%s_%s_%s_%s_%s_summary.json" % (
            self.model_name, gpu_type, self.mean_input_tokens, self.mean_output_tokens, self.num_concurrent_requests)

    def gen_test_commands(self):
        test_commands = [
            'export OPENAI_API_KEY=' + openai_api_key,
            'export OPENAI_API_BASE=' + openai_api_base,
            'source ' + env_activate_path,
            'cd ' + llmperf_path,
            'python3 token_benchmark_ray.py --model ' + self.model_name + ' --mean-input-tokens ' + \
            self.mean_input_tokens + ' --stddev-input-tokens ' + self.stddev_input_tokens + ' --mean-output-tokens ' + \
            self.mean_output_tokens + ' --stddev-output-tokens ' + self.stddev_output_tokens + \
            ' --max-num-completed-requests ' + self.max_num_completed_requests + ' --timeout ' + self.timeout + \
            ' --num-concurrent-requests ' + self.num_concurrent_requests + ' --results-dir "result_outputs" --llm-api openai'
        ]
        return test_commands

    def run_model_test(self):

        # 创建SSH客户端实例
        ssh_client = SSHClient(ip, port, username, password)

        # 创建自动化测试实例
        test_commands = self.gen_test_commands()

        automated_test = AutomatedTest(ssh_client, test_commands, self.remote_path, self.local_path)

        return automated_test.run_test()

    def save_test_reulsts(self):
        results_file = f'./results/apiperf-test-results.xlsx'
        with open(self.local_path, 'r', encoding='utf-8') as file:

            # 提取测试结果中的指标
            data = json.load(file)
            model_name = data["model"]
            concurrent = data["num_concurrent_requests"]
            itl = data["results_inter_token_latency_s_mean"]
            ttft = data["results_ttft_s_mean"]
            etel = data["results_end_to_end_latency_s_mean"]
            rottps = data["results_request_output_throughput_token_per_s_mean"]
            nitm = data["results_number_input_tokens_mean"]
            notm = data["results_number_output_tokens_mean"]

            # 创建excel对象写入测试数据
            if not os.path.exists(results_file):
                # 文件不存在则创建一个新的工作薄和工作表
                wb = Workbook()
                ws = wb.active
                ws.title = model_name
                ws.append(["模型", "实例规格", "并发", "ITL", "TTFT", "端到端时延", "输出token吞吐量", "输入token数",
                           "输出token数"])
                # 追加数据
                ws.append([model_name, gpu_type, concurrent, itl, ttft, etel, rottps, nitm, notm])
                wb.save(results_file)
                print(f"创建文件：{results_file}")
            else:
                # 文件存在则加载工作薄
                wb = load_workbook(results_file)

                # 检查工作表是否存在，不存在则创建
                if model_name in wb.sheetnames:
                    ws = wb[model_name]
                else:
                    ws = wb.create_sheet(model_name)
                    ws.append(
                        ["模型", "实例规格", "并发", "ITL", "TTFT", "端到端时延", "输出token吞吐量", "输入token数",
                         "输出token数"])
                # 追加数据
                ws.append([model_name, gpu_type, concurrent, itl, ttft, etel, rottps, nitm, notm])

                # 保存文件
                wb.save(results_file)
                print(f"写入数据到文件{results_file}成功。")


class LlmCorrectnessTest():
    def __init__(self, *args, **kwargs):

        # 定义默认值
        defaults = {
            'model_name': 'baichuan2-13B-Chat',
            'max_num_completed_requests': '150',
            'timeout': '600',
            'num_concurrent_requests': '10'
        }

        # 更新默认值
        defaults.update(kwargs)

        # 参数赋值给实例变量
        self.model_name = defaults['model_name']
        self.max_num_completed_requests = defaults['max_num_completed_requests']
        self.timeout = defaults['timeout']
        self.num_concurrent_requests = defaults['num_concurrent_requests']

        # 远程测试结果文件路径
        self.remote_path = "%s/result_outputs/%s_correctness_summary.json" % (
            llmperf_path, self.model_name)
        # 下载到本地的文件路径
        if not os.path.exists('./results'):
            os.makedirs('./results')
        self.local_path = "./results/%s_correctness_summary.json" % (
            self.model_name)

    def gen_test_commands(self):
        test_commands = [
            'export OPENAI_API_KEY=' + openai_api_key,
            'export OPENAI_API_BASE=' + openai_api_base,
            'source ' + env_activate_path,
            'cd ' + llmperf_path,
            'python3 llm_correctness.py --model ' + self.model_name + ' --max-num-completed-requests ' + \
            self.max_num_completed_requests + ' --timeout ' + self.timeout + ' --num-concurrent-requests ' + \
            self.num_concurrent_requests + ' --results-dir "result_outputs"'
        ]
        return test_commands

    def run_model_test(self):

        # 创建SSH客户端实例
        ssh_client = SSHClient(ip, port, username, password)

        # 创建自动化测试实例
        test_commands = self.gen_test_commands()

        automated_test = AutomatedTest(ssh_client, test_commands, self.remote_path, self.local_path)

        return automated_test.run_test()

    def save_test_reulsts(self):
        results_file = f'./results/correctness-test-results.xlsx'
        with open(self.local_path, 'r', encoding='utf-8') as file:

            # 提取测试结果中的指标
            data = json.load(file)
            model_name = data["model"]
            number_errors = data["number_errors"]
            num_mismatched_requests = data["num_mismatched_requests"]
            error_rate = data["error_rate"]
            mismatch_rate = data["mismatch_rate"]
            num_completed_requests = data["num_completed_requests"]
            num_non_errored_requests = data["num_non_errored_requests"]
            num_concurrent_requests = data["num_concurrent_requests"]

            # 创建excel对象写入测试数据
            if not os.path.exists(results_file):
                # 文件不存在则创建一个新的工作薄和工作表
                wb = Workbook()
                ws = wb.active
                ws.title = model_name
                ws.append(
                    ["模型", "错误数", "不匹配的请求数", "错误比例", "不匹配比例", "完成的请求数", "非错误的请求数",
                     "并发的请求数"])
                # 追加数据
                ws.append([model_name, number_errors, num_mismatched_requests, error_rate, mismatch_rate,
                           num_completed_requests, num_non_errored_requests, num_concurrent_requests])
                wb.save(results_file)
                print(f"创建文件：{results_file}")
            else:
                # 文件存在则加载工作薄
                wb = load_workbook(results_file)

                # 检查工作表是否存在，不存在则创建
                if model_name in wb.sheetnames:
                    ws = wb[model_name]
                else:
                    ws = wb.create_sheet(model_name)
                    ws.append(
                        ["模型", "错误数", "不匹配的请求数", "错误比例", "不匹配比例", "完成的请求数", "非错误的请求数",
                         "并发的请求数"])
                # 追加数据
                ws.append([model_name, number_errors, num_mismatched_requests, error_rate, mismatch_rate,
                           num_completed_requests, num_non_errored_requests, num_concurrent_requests])

                # 保存文件
                wb.save(results_file)
                print(f"写入数据到文件{results_file}成功。")
